import dataclasses
import datetime
import itertools
import logging
from typing import List, Optional
from absl import app

from coin.base.datetime_util import (
    to_datetime, to_timestamp_int)
from coin.proto.coin_market_query_pb2 import ProductTicker
from coin.proto.coin_strategy_pb2 import StrategyLog
from coin.support.proto_log.logic.util import (read_strat_log, StratInfo)


@dataclasses.dataclass
class ReserveChange:
  ts: int
  qty: float
  price: ProductTicker


class OnLog(object):
  def __init__(self, symbol: str, hypo_reserve: float, start_ts: int):
    self._symbol = symbol
    self._base = symbol.split('-')[0]
    self._last_reserve: Optional[float] = None
    self._last_ticker: Optional[ProductTicker] = None
    hypo_reserve = ReserveChange(start_ts, hypo_reserve, None)
    self._reserve_changes: List[ReserveChange] = [hypo_reserve]

  def _handle_reserve(self, ts, pb: StrategyLog):
    for each_balance in pb.reserve.each_balance:
      if each_balance.currency == self._base:
        total = each_balance.total
        if self._last_reserve is None or self._last_reserve != total:
          logging.info(f'reserve changed: {to_datetime(ts)} from {self._last_reserve} to {total}')
          reserve_change = ReserveChange(ts, total, self._last_ticker)
          self._reserve_changes.append(reserve_change)
        self._last_reserve = total

  def _handle_ticker(self, ts, pb: StrategyLog):
    for each_ticker in pb.ticker.each_ticker:
      if each_ticker.symbol == self._symbol:
        each_ticker.exchange_timestamp = ts
        hypo_reserve = self._reserve_changes[0]
        last_reserve_change = self._reserve_changes[-1]
        if hypo_reserve.price is None:
          hypo_reserve.price = each_ticker
        last_reserve_change.price = each_ticker
        self._last_ticker = each_ticker
  
  def on_log(self, ts: int, log: bytes):
    pb = StrategyLog()
    pb.ParseFromString(log)
    if pb.type == StrategyLog.RESERVE:
      self._handle_reserve(ts, pb)
    elif pb.type == StrategyLog.TICKER:
      self._handle_ticker(ts, pb)

  def get_pnl_diff(self) -> float:
    pnl_diff = 0
    hypo_reserve = self._reserve_changes[0]
    for x, y in itertools.pairwise(self._reserve_changes):
      if x.price is not None and y.price is not None:
        pnl_diff += (y.price.last - x.price.last) * (y.qty - hypo_reserve.qty)
    return pnl_diff


def main(_):
  strategy_name = 'vmm_ccx_delta_hedge'
  machine = 'strategy-21.ap-southeast-1'
  root_dir = '/remote/iosg/strat-1/buckets/log.raw.coin/live/strat_proto_log'
  start_time = datetime.datetime(2023, 1, 20)
  end_time = datetime.datetime(2023, 1, 21)
  pnl_diff_calc = OnLog(
      symbol='CCX-USDT', hypo_reserve=2*10**7, start_ts=to_timestamp_int(start_time))
  strat_info = StratInfo(strategy_name=strategy_name,
                         machine=machine,
                         trading_date=None)
  read_strat_log(strat_info=strat_info,
                 start_time=start_time,
                 end_time=end_time,
                 callback=pnl_diff_calc.on_log,
                 root_dir=root_dir)
  print(pnl_diff_calc.get_pnl_diff())

if __name__ == '__main__':
  
  logging.basicConfig(level='DEBUG', format='%(levelname)8s %(asctime)s %(name)s] %(message)s')
  app.run(main)
